import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
import datetime
from kneed import KneeLocator


def getFactor(code, data):
    data = data.sort_values(by='trade_date', ascending=True)
    data['factor'] = (data['close'] - data['high'].shift(20).max()) / (
                data['high'].shift(20).max() - data['low'].shift(20).min())

    #plt.plot(data['trade_date'], data['factor'])
    #plt.show()

    kneedle_con_inc = KneeLocator(data.index,
                                  data['close'],
                                  curve='concave',
                                  direction='decreasing',
                                  online=False)

    if kneedle_con_inc.knee != None:
        print(code, kneedle_con_inc.knee, kneedle_con_inc.knee_y)

    data = data.sort_values(by='factor')
    #print(data)

    return data

def getRecentDay(trade_date):
    days = []
    days.clear()
    days.append(trade_date)
    dtend = datetime.datetime.strptime(trade_date, '%Y%m%d').date()
    for i in range(1,10):
        daytemp = dtend - datetime.timedelta(days=i)
        day = daytemp.strftime('%Y%m%d')
        days.append(day)

    return days

def deal(code, days, data):
    factordata = getFactor(code, data)
    fivedays = factordata.iloc[0:10, 1].values
    jiaoji = list(set(days).intersection(set(fivedays)))
    if len(jiaoji) > 0:
        print(code)

    # print(days)
    # print(fivedays)
    # for day in fivedays:
    #     if day in days:
    #         print(code)
    #         break

def getall(trade_day):
    ts.set_token('6f5fb1b9f9e12bf2b8df3ea881697964040b76e20ff084ceb3efc4b3')
    pro = ts.pro_api()

    days = getRecentDay(trade_day)
    # 选取主板股票列表，排除创业板，科创板
    stocklist = pro.stock_basic(exchange='', list_status='L', market='主板,创业板', fields='ts_code,symbol,name')
    print(stocklist)
    num = 0
    codestr = ''
    for stock in stocklist.itertuples():
        ts_code = getattr(stock, 'ts_code')
        codestr += ts_code
        codestr += ','
        num = num + 1
        if num == 50:
            # 批量获取行情数据
            #print(codestr)
            datas = pro.daily(ts_code=codestr, start_date='20210801', end_date='20211220')
            codes = codestr.split(',')
            #print(codes)
            for code in codes:
                if code != '':
                    data = datas[datas['ts_code'] == code]
                    deal(code, days, data)
                    break

            codestr = ''
            num = 0
            # # 行情数据过滤
            # for df in datas.itertuples():
            #     ts_code = getattr(df, 'ts_code')
            #     open = getattr(df, 'open')
            #     close = getattr(df, 'close')
            #     trade_date = getattr(df, 'trade_date')


if __name__ == '__main__':
    trade_day = '20211220'
    getall(trade_day)




